// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/benchmark.h>
#include <melon/init/init.h>

#include <turbo/flags/flag.h>

#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/vector/vector_fuzzer.h>

TURBO_FLAG(int64_t,fuzzer_seed, 99887766, "Seed for random input dataset generator");

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::test;

namespace {
const std::string_view kColName{"a"};

class SimpleCastBenchmark : public functions::test::FunctionBenchmarkBase {
 public:
  explicit SimpleCastBenchmark() : FunctionBenchmarkBase() {}

  RowVectorPtr make_row_vector(const TypePtr& type, vector_size_t size) {
    VectorFuzzer::Options opts;
    opts.vectorSize = size;
    opts.nullRatio = 0;
    VectorFuzzer fuzzer(std::move(opts), pool(), FLAGS_fuzzer_seed);
    VectorPtr input = fuzzer.fuzzFlat(type); // Col a
    return vectorMaker_.row_vector(
        std::vector<std::string>{std::string(kColName)},
        std::vector<VectorPtr>{std::move(input)});
  }

  static constexpr auto kNumSmallRuns = 100;
  static constexpr auto kNumMediumRuns = 1000;
  static constexpr auto kNumLargeRuns = 10'000;

  void runSmall(const TypePtr& inputType, const TypePtr& outputType) {
    run(inputType, outputType, kNumSmallRuns);
  }

  void runMedium(const TypePtr& inputType, const TypePtr& outputType) {
    run(inputType, outputType, kNumMediumRuns);
  }

  void runLarge(const TypePtr& inputType, const TypePtr& outputType) {
    run(inputType, outputType, kNumLargeRuns);
  }

  size_t
  run(const TypePtr& inputType, const TypePtr& outputType, size_t batchSize) {
    melon::BenchmarkSuspender suspender;
    auto input = make_row_vector(inputType, batchSize);
    auto castInput =
        std::make_shared<kumo::pollux::core::FieldAccessTypedExpr>(
            inputType, std::string(kColName));
    std::vector<kumo::pollux::core::TypedExprPtr> expr{
        std::make_shared<kumo::pollux::core::CastTypedExpr>(
            outputType, castInput, false)};
    exec::ExprSet exprSet(expr, &execCtx_);
    SelectivityVector rows(input->size());
    VectorPtr result;
    suspender.dismiss();

    size_t count = 0;
    for (auto i = 0; i < 1000; i++) {
      evaluate(exprSet, input, rows, result);
      count += result->size();
    }
    return count;
  }
};

TypePtr buildStructType(
    std::function<std::string(int)>&& nameGenerator,
    const TypePtr& fieldType,
    size_t numChildren) {
  std::vector<std::string> names;
  std::vector<TypePtr> types;
  for (int i = 0; i < numChildren; i++) {
    names.push_back(nameGenerator(i));
    types.push_back(fieldType);
  }
  return ROW(std::move(names), std::move(types));
}

std::unique_ptr<SimpleCastBenchmark> benchmark;

BENCHMARK(castTimestampDateSmall) {
  benchmark->setAdjustTimestampToTimezone("false");
  benchmark->runSmall(TIMESTAMP(), DATE());
}

BENCHMARK(castTimestampDateMedium) {
  benchmark->setAdjustTimestampToTimezone("false");
  benchmark->runMedium(TIMESTAMP(), DATE());
}

BENCHMARK(castTimestampDateLarge) {
  benchmark->setAdjustTimestampToTimezone("false");
  benchmark->runLarge(TIMESTAMP(), DATE());
}

BENCHMARK(castTimestampDateAdjustTimeZoneSmall) {
  benchmark->setTimezone("America/Los_Angeles");
  benchmark->setAdjustTimestampToTimezone("true");
  benchmark->runSmall(TIMESTAMP(), DATE());
}

BENCHMARK(castTimestampDateAdjustTimeZoneMedium) {
  benchmark->setTimezone("America/Los_Angeles");
  benchmark->setAdjustTimestampToTimezone("true");
  benchmark->runMedium(TIMESTAMP(), DATE());
}

BENCHMARK(castTimestampDateAdjustTimeZoneLarge) {
  benchmark->setTimezone("America/Los_Angeles");
  benchmark->setAdjustTimestampToTimezone("true");
  benchmark->runLarge(TIMESTAMP(), DATE());
}

BENCHMARK(castStructFewFieldsRenameSmall) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 3);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, INTEGER(), 3);
  suspender.dismiss();

  benchmark->runSmall(oldType, newType);
}

BENCHMARK(castStructFewFieldsRenameMedium) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 3);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, INTEGER(), 3);
  suspender.dismiss();

  benchmark->runMedium(oldType, newType);
}

BENCHMARK(castStructFewFieldsRenameLarge) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 3);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, INTEGER(), 3);
  suspender.dismiss();

  benchmark->runLarge(oldType, newType);
}

BENCHMARK(castStructManyFieldsRenameSmall) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 1000);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, INTEGER(), 1000);
  suspender.dismiss();

  benchmark->runSmall(oldType, newType);
}

BENCHMARK(castStructManyFieldsRenameMedium) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 1000);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, INTEGER(), 1000);
  suspender.dismiss();

  benchmark->runMedium(oldType, newType);
}

BENCHMARK(castStructManyFieldsRenameLarge) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 1000);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, INTEGER(), 1000);
  suspender.dismiss();

  benchmark->runLarge(oldType, newType);
}

BENCHMARK(castStructFewFieldsNestedCastSmall) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 3);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, BIGINT(), 3);
  suspender.dismiss();

  benchmark->runSmall(oldType, newType);
}

BENCHMARK(castStructFewFieldsNestedCastMedium) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 3);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, BIGINT(), 3);
  suspender.dismiss();

  benchmark->runMedium(oldType, newType);
}

BENCHMARK(castStructFewFieldsNestedCastLarge) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 3);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, BIGINT(), 3);
  suspender.dismiss();

  benchmark->runLarge(oldType, newType);
}

BENCHMARK(castStructManyFieldsNestedCastSmall) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 1000);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, BIGINT(), 1000);
  suspender.dismiss();

  benchmark->runSmall(oldType, newType);
}

BENCHMARK(castStructManyFieldsNestedCastMedium) {
  melon::BenchmarkSuspender suspender;
  auto oldType = buildStructType([](int) { return ""; }, INTEGER(), 1000);
  auto newType = buildStructType(
      [](int i) { return fmt::format("col{}", i); }, BIGINT(), 1000);
  suspender.dismiss();

  benchmark->runMedium(oldType, newType);
}

// castStructManyFieldsNestedCastLarge is skipped because it takes too long to
// run.

} // namespace

int main(int argc, char* argv[]) {
  melon::Init init{&argc, &argv};
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  memory::MemoryManager::initialize({});
  benchmark = std::make_unique<SimpleCastBenchmark>();
  melon::runBenchmarks();
  benchmark.reset();
  return 0;
}
